import os
import re
import linecache
import numpy as np
import numpy.fft as nf
import matplotlib.pyplot as plt


def get_all_data_files(folder):
    """
    获得文件夹中所有的数据文件
    """
    files = os.listdir(folder)
    files.sort()
    for file in files:
        if file.endswith(".dat"):
            yield rf"{folder}/{file}"


def get_data(file):
    """
    读取文件中的数据
    """
    data = np.loadtxt(file, skiprows=9)  # 跳过前9行数据
    # 到底是否需要乘以比例系数，需要确认！！！
    # scale = float(re.split(r'[\t\n]', linecache.getline(file, 2))[1])  # 读取信号的比例值
    # data[:, 1] = data[:, 1] * scale  # 将第二列数据乘以比例值
    return data


def get_frequency(file):
    """
    获得文件的记录频率

    :param file: 包含上级文件夹的文件名
    :return: 记录频率
    """
    return float(file.split('/')[1].split('Hz')[0])  # 从文件名中读取激振器的频率值


def get_f_A(t_y):
    """
    对数据进行fft变换，得到频域图需要的数据
    :param t_y: 需要分析的数据
    :return: 频域图的频率序列和幅值序列
    """
    t = t_y[:, 0]  # 第一列数据为时间序列
    y = t_y[:, 1]  # 第二列数据为函数序列
    Y = nf.fft(y)
    N = y.size  # 采样点数
    delta_t = t[1] - t[0]  # 采样间隔
    f_a = nf.fftfreq(y.size, delta_t)  # 频域的横坐标序列
    mask = np.where(f_a > 0)  # 只分析频率大于零的部分（小于零的部分与之完全对称）
    f = f_a[mask]  # 取出频率序列中大于零的部分
    A = 2 * abs(Y[mask]) / N  # 频谱图的纵坐标变换为实际振幅
    return np.array([f, A]).transpose()  # 返回频率序列和幅值序列


def draw_figures(t_y, f_A, meas_freq, file, file_name_mark='', clear_low_frequency=True):
    """
    绘制频域图
    :param t_y: 时域图的时间序列和函数序列
    :param f_A: 频域图的频率序列和幅值序列
    :param meas_freq: 测量的振动频率
    :param file: 文件名
    :param file_name_mark: 文件名中的标注
    :param clear_low_frequency: 是否清除低频干扰
    :return: None
    """
    t = t_y[:, 0]
    y = t_y[:, 1]
    f = f_A[:, 0]
    A = f_A[:, 1]
    if clear_low_frequency:
        mask = np.where(f > meas_freq * 0.1)
        f = f[mask]
        A = A[mask]
    amplitude = np.max(A)  # 最大振幅值
    sign_freq = f[np.argmax(A)]  # 找到频谱图山振幅最大时对应的频率值
    print(f"\n分量最大的频率分量为：{sign_freq:.2f} Hz, 对应的幅值为：{amplitude:.2f} V")
    devi_freq = sign_freq - meas_freq  # 绝对误差
    real_devi = devi_freq / meas_freq  # 相对误差
    print(f"偏差为：{devi_freq:.2f} Hz（{real_devi:.2%}）")
    text = file + "\n" + f'f = {sign_freq:.2f} Hz, A = {amplitude:.3e} V, ' \
                         f'error = {devi_freq:.2f} Hz, {real_devi:.2%}'
    fig, axes = plt.subplots(1, 2)  # 画两个子图，分别为时域图和频域图
    axes[0].plot(t, y)
    axes[0].set_xlabel("time (s)")
    axes[0].set_ylabel("signal")
    plt.suptitle(text)  # 将误差作为图的标题，显示在图的上方

    axes[1].plot(f, A)
    axes[1].set_xlabel("frequency (Hz)")
    axes[1].set_ylabel("$A$")
    plt.tight_layout(rect=[0, 0, 1, 1])  # 设置子图布局，避免重叠
    plt.savefig(f'{file}_{file_name_mark}.pdf')  # 将分析结果保存在以频率信息命名的pdf图片中
    plt.close()


def filter_freq(t_y, freq_list=None, band_width=1):
    """
    对频率列表进行滤波，可以指定滤波的带宽
    :param t_y: 时域图的时间序列和函数序列
    :param freq_list: 频率列表，默认为[0, 50]，表示删除0 Hz和50 Hz左右的分量
    :param band_width: 滤波带宽
    :return: 过滤后的时域信号和频域信号
    """
    t = t_y[:, 0]  # 第一列数据为时间序列
    y = t_y[:, 1]  # 第二列数据为函数序列
    Y = nf.fft(y)
    N = y.size  # 采样点数
    delta_t = t[1] - t[0]  # 采样间隔
    f = nf.fftfreq(y.size, delta_t)  # 频域的横坐标序列
    mask = np.where(f > 0)  # 只分析频率大于零的部分（小于零的部分与之完全对称）
    A = 2 * abs(Y[mask]) / N  # 频谱图的纵坐标变换为实际振幅
    filtered_y = []

    if freq_list is None:
        freq_list = [0, 50]

    for freq in freq_list:
        lower_freq = freq - band_width/2
        upper_freq = freq + band_width/2
        filtered_Y = Y * ((abs(f) < lower_freq) | (abs(f) > upper_freq))   # 通过布尔运算，将频率在指定范围内的分量置零
        filtered_y = nf.ifft(filtered_Y)
        y = filtered_y.real
        Y = filtered_y

    filtered_A = 2 * abs(filtered_y[mask]) / N  # 频域的纵坐标序列
    filtered_f_A = np.array([f[mask], filtered_A]).transpose()  # 返回频率序列和幅值序列
    filtered_t_y = np.array([t, y]).transpose()  # 返回时域序列和函数序列

    return filtered_t_y, filtered_f_A


def write_data(data, file_name, digit=4):
    """
    将数据写入文件
    :param data: 需要写入的数据
    :param file_name: 文件名
    :return: None
    """
    with open(file_name, 'w') as f:
        for i in data:
            f.write(f"{i[0]:.{digit}f}\t{i[1]:.{digit}f}\n")


def noise_reduce(raw_data, N=3, Wn=0.01, output='ba'):
    """
    降噪函数，使用Butterworth滤波器
    :param raw_data:
    :param N:过滤的阶数，可以修改
    :param Wn:截断频率，可以修改
    :param output:输出的格式，可以修改
    :return:
    """
    import scipy.signal as signal
    # B, A = signal.firwin(N, Wn, output=output)
    B, A = signal.butter(N, Wn, output=output)
    # raw_data[:, 1] = signal.filtfilt(B, A, raw_data[:, 1])    # 这样写会破坏原始数据raw_data
    filtered_data = signal.filtfilt(B, A, raw_data[:, 1])
    return np.array([raw_data[:, 0], filtered_data]).transpose()


def draw_y(data, file_name, file_name_mark):
    """
    绘制函数序列
    :param data: 时域图的时间序列和函数序列
    :return: None
    """
    t = data[:, 0]  # 第一列数据为时间序列
    y = data[:, 1]  # 第二列数据为函数序列
    plt.plot(t, y)
    plt.xlabel("time (s)")
    plt.ylabel("signal")
    plt.savefig(f"{file_name}_{file_name_mark}.pdf")
    plt.close()


if __name__ == '__main__':
    directory = 'Data'
    files = get_all_data_files(directory)
    for file in files:
        t_y = get_data(file)
        draw_y(t_y, file, 'raw')
        # t_y_nr = noise_reduce(t_y, N=3, Wn=0.01, output='ba')
        # draw_y(t_y_nr, file, 'noise_reduce')
        meas_freq = get_frequency(file)
        f_A = get_f_A(t_y)
        draw_figures(t_y, f_A, meas_freq, file, "time_domain_and_frequency_domain")
        # # write_data(data, f'{i}_t_y.txt')
        # # write_data(f_and_A, f'{i}_f_A.txt')
        # filtered_t_y, filtered_f_A = filter_freq(t_y, freq_list=[0, 50, 56], band_width=2)
        # draw_figures(filtered_t_y, filtered_f_A, meas_freq, file, file_name_mark='filtered')


